# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Created by Tianheng Cheng(tianhengcheng@gmail.com), Yang Zhao
# ------------------------------------------------------------------------------

import math

import torch
import numpy as np

from ..utils.transforms import transform_preds


def get_preds(scores):
    """
    get predictions from score maps in torch Tensor
    :arg scores: [b,n,w,h]
    return type: torch.LongTensor with size=[b,n,2], ordered by [w,h]
    """
    assert scores.dim() == 4, 'Score maps should be 4-dim'
    maxval, idx = torch.max(scores.view(scores.size(0), scores.size(1), -1), 2)

    maxval = maxval.view(scores.size(0), scores.size(1), 1)
    idx = idx.view(scores.size(0), scores.size(1), 1) + 1

    preds = idx.repeat(1, 1, 2).float()

    preds[:, :, 1] = (preds[:, :, 1] - 1) % scores.size(3) + 1
    preds[:, :, 0] = torch.floor((preds[:, :, 0] - 1) / scores.size(3)) + 1

    pred_mask = maxval.gt(0).repeat(1, 1, 2).float()
    preds *= pred_mask
    return preds


def compute_nme(preds, meta):
    '''

    :param preds: [b,n,2] where points ordered by [w,h]
    :param meta: meta['pts']=[b,n,2] where points ordered by [w,h]
    :return:
    '''
    targets = meta['pts']
    preds = preds.numpy()
    target = targets.cpu().numpy()

    N = preds.shape[0]
    L = preds.shape[1]
    rmse = np.zeros(N)

    for i in range(N):
        pts_pred = preds[i,]
        pts_gt = target[i,]
        if L == 19:  # aflw
            interocular = meta['box_size'][i]
        elif L == 29:  # cofw
            interocular = np.linalg.norm(pts_gt[8,] - pts_gt[9,])
        elif L == 68:  # 300w
            # interocular
            interocular = np.linalg.norm(pts_gt[36,] - pts_gt[45,])
        elif L == 98:
            interocular = np.linalg.norm(pts_gt[60,] - pts_gt[72,])
        elif L == 2:
            interocular = np.linalg.norm(pts_gt[0,] - pts_gt[1,])
            if interocular < 1e-1:
                # bad label encounted
                # print('pts_gt=', pts_gt)
                # print('pts_pred=', pts_pred)
                # print('L=', L)
                interocular = 300
        else:
            raise ValueError('Number of landmarks is wrong')
        rmse[i] = np.sum(np.linalg.norm(pts_pred - pts_gt, axis=1)) / (interocular * L)

    return rmse


def decode_preds(output, center, scale, res):
    '''
    :param output: heatmap with size=[b,npts,w,h]
    :param center: [w,h]
    :param scale:
    :param res: [w,h]
    :return: preds: size=[b,c,2] where dim2=[w,h]
    '''
    coords = get_preds(output)  # float type [b,n,2] where dim2=[w,h]

    coords = coords.cpu()
    # pose-processing
    for n in range(coords.size(0)):
        for p in range(coords.size(1)):
            hm = output[n][p]
            px = int(math.floor(coords[n][p][1]))
            py = int(math.floor(coords[n][p][0]))
            # debug
            # print('error pt=', px, ',', py)
            if (px > 1) and (px < res[1]) and (py > 1) and (py < res[0]):
                diff = torch.Tensor([hm[py - 1][px] - hm[py - 1][px - 2], hm[py][px - 1] - hm[py - 2][px - 1]])
                coords[n][p] += diff.sign() * .25
    coords += 0.5
    preds = coords.clone()

    # Transform back
    for i in range(coords.size(0)):
        preds[i] = transform_preds(coords[i], center[i], scale[i], res)

    if preds.dim() < 3:
        preds = preds.view(1, preds.size())

    return preds
